import Utils (tok)
import Data.Char (ord, isDigit)

hash :: String -> Int
hash = 
    go 0 
    where
        go cur []     = cur
        go cur (c:cs) = go cur' cs
        where
            cur' = ((cur + (ord c)) * 17) `rem` 256

type Lens      = (LensLabel, FocLength)
type LensLabel = String
type FocLength = Int

data Command   = Remove CmdRemove | Add CmdAdd deriving (Show)
type CmdRemove = (BoxNb, LensLabel)
type CmdAdd    = (BoxNb, LensLabel, FocLength)

type BoxNb     = Int
type Boxchain  = [[Lens]]

parseCommand :: String -> Command
parseCommand s
    | '-' `elem` s = Remove (box, label)
    | '=' `elem` s = Add    (box, label, focLen)
    where
        label  = takeWhile (not . (`elem` "=-")) s
        box    = hash label
        focLen = read $ dropWhile (not . isDigit) $ s

removeLens :: Boxchain -> BoxNb -> LensLabel -> Boxchain
removeLens boxes i lab = 
    (take i boxes) ++ [box'] ++ (drop (i + 1) boxes)
    where
        box' = filter ((/=lab) . fst) (boxes !! i)

addLens :: Boxchain -> BoxNb -> Lens -> Boxchain
addLens boxes i (lab, foc) = 
    (take i boxes) ++ [box'] ++ (drop (i + 1) boxes)
    where
        box   = boxes !! i
        isNew = lab `elem` (map fst box)
        box'  = 
            if not isNew then
                box ++ [(lab, foc)]
            else
                left ++ [(lab, foc)] ++ right
        where
            left  = takeWhile ((/=lab) . fst) box
            right = drop 1 . dropWhile ((/=lab) . fst) $ box

execCommand :: Boxchain -> Command -> Boxchain
execCommand b (Remove (nb, lab))   = removeLens b nb lab
execCommand b (Add (nb, lab, foc)) = addLens b nb (lab, foc)

foc :: Boxchain -> Int
foc b = sum
      . map (uncurry (*))
      . zip [1..]
      . map (foldl (+) 0)
      . map (map (uncurry (*)))
      . map (zip [1..])
      . map (map (\(_, fl) -> fl))
      $ b

main = do
    contents <- readFile "input.txt"
    let input = tok "," $ head . lines $ contents

    print $ sum . map hash $ input

    let commands = map parseCommand input
    let empty = replicate 256 []

    let final = foldl execCommand empty commands
    print $ foc $ final